package org.theSeed.web.filter;

import cn.hutool.core.lang.UUID;
import cn.hutool.core.util.ObjectUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;
import org.springframework.core.Ordered;
import org.theSeed.web.sysEnum.FilterOrder;
import org.theSeed.web.utils.ServletUtil;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.IOException;
import java.util.*;

/**
 * 日志过滤器
 */
@Slf4j
public class TraceFilter implements Filter, Ordered {
    public TraceFilter() {
    }

    public TraceFilter(Integer orderNum) {
        this.orderNum = orderNum;
    }

    public TraceFilter(String headerName) {
        this.headerName = headerName;
    }

    public TraceFilter(Integer orderNum,String headerName) {
        this.orderNum = orderNum;
        this.headerName = headerName;
    }

    protected Integer orderNum = FilterOrder.TRACE.getCode();

    /**
     * 头名称
     */
    protected String headerName = "TRACE_ID";

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        String traceId = ServletUtil.getHeaderValue((HttpServletRequest) servletRequest, headerName);
        if (ObjectUtil.isEmpty(traceId)){
            HttpServletRequest httpRequest = (HttpServletRequest) servletRequest;
            // 创建自定义的请求包装器
            CustomHeaderRequestWrapper wrappedRequest = new CustomHeaderRequestWrapper(httpRequest);
            servletRequest = wrappedRequest;
            // 添加自定义请求头
            wrappedRequest.addHeader(headerName, UUID.randomUUID().toString());

        }
        MDC.put(headerName,traceId);
        filterChain.doFilter(servletRequest, servletResponse);
    }

    @Override
    public int getOrder() {
        return orderNum;
    }

    // 自定义请求包装器
    public class CustomHeaderRequestWrapper extends HttpServletRequestWrapper {

        private final Map<String, String> customHeaders = new HashMap<>();

        public CustomHeaderRequestWrapper(HttpServletRequest request) {
            super(request);
        }

        public void addHeader(String name, String value) {
            customHeaders.put(name, value);
        }

        @Override
        public String getHeader(String name) {
            // 先检查自定义头
            String customValue = customHeaders.get(name);
            if (customValue != null) {
                return customValue;
            }
            // 否则返回原始头的值
            return super.getHeader(name);
        }

        @Override
        public Enumeration<String> getHeaders(String name) {
            // 处理自定义头和多值头
            if (customHeaders.containsKey(name)) {
                return Collections.enumeration(Collections.singletonList(customHeaders.get(name)));
            }
            return super.getHeaders(name);
        }

        @Override
        public Enumeration<String> getHeaderNames() {
            // 合并自定义头和原始头的名称
            Set<String> allHeaderNames = new HashSet<>();
            allHeaderNames.addAll(customHeaders.keySet());
            Enumeration<String> originalHeaderNames = super.getHeaderNames();
            while (originalHeaderNames.hasMoreElements()) {
                allHeaderNames.add(originalHeaderNames.nextElement());
            }
            return Collections.enumeration(allHeaderNames);
        }

        @Override
        public long getDateHeader(String name) {
            String value = getHeader(name);
            if (value == null) {
                return -1L;
            }
            try {
                return Long.parseLong(value);
            } catch (NumberFormatException e) {
                throw new IllegalArgumentException("Invalid date header: " + name);
            }
        }
    }
}
